#include "VideoEncoder.h"
#include "common.h"

CVideoEncoder::CVideoEncoder(EncodeDataNotify *pNotify)
	: m_pNotify(pNotify)
	, m_hEncodeHandle(NULL)
	, m_bExit(false)
	, m_pH264Buf(NULL)
	, m_pFrameRGB(NULL)
	, m_pRGBBuff(NULL)
	, m_pFrameYUV(NULL)
	, m_pYUVBuff(NULL)
	, m_pRGB_YUV_CTX(NULL)
	, m_nSrcWidth(0)
	, m_nSrcHeight(0)
	, m_bEncode(false)
{
	StartThread();
}

CVideoEncoder::~CVideoEncoder(void)
{
	m_bExit = true;

	WaitForStop();

	CloseEncoder();

	while (m_lstRGBData.size())
	{
		RGBDataBuff *pRGB = NULL;

		{
			CAutoLock autoLock(&m_lockRGBData);
			pRGB = (RGBDataBuff *)m_lstRGBData.front();
			m_lstRGBData.pop_front();
		}

		if (pRGB)
			ClearRGBBuf(pRGB);
	}

	if (m_pH264Buf)
	{
		delete [] m_pH264Buf;
		m_pH264Buf = NULL;
	}
}

void CVideoEncoder::PutRGBData(const char *pData, const int &nLen, int nWidth, int nHeight, unsigned long uiTimestamp)
{
	if (m_bEncode)
	{
		RGBDataBuff *pRGBData = new RGBDataBuff;
		if (pRGBData)
		{
			pRGBData->pData = new char[nLen];
			pRGBData->nDataLen = nLen;
			pRGBData->nWidth = nWidth;
			pRGBData->nHeigth = nHeight;
			pRGBData->ulTimestamp = uiTimestamp;

			if (pRGBData->pData)
			{
				memcpy(pRGBData->pData, pData, nLen);

				CAutoLock autoLock(&m_lockRGBData);
				m_lstRGBData.push_back(pRGBData);
			}
			else
			{
				delete pRGBData;
				pRGBData = NULL;
			}
		}
	}	
}

void CVideoEncoder::StartEncode(void)
{
	m_bEncode = true;
}

void CVideoEncoder::StopEncode(void)
{
	m_bEncode = false;

	while (m_lstRGBData.size())
	{
		RGBDataBuff *pRGB = NULL;

		{
			CAutoLock autoLock(&m_lockRGBData);
			pRGB = (RGBDataBuff *)m_lstRGBData.front();
			m_lstRGBData.pop_front();
		}

		if (pRGB)
			ClearRGBBuf(pRGB);
	}
}

void CVideoEncoder::ClearRGBBuf(RGBDataBuff *pRGB)
{
	if (pRGB)
	{
		if (pRGB->pData)
		{
			delete [] pRGB->pData;
			pRGB->pData = NULL;
		}

		delete pRGB;
		pRGB = NULL;
	}
}

bool CVideoEncoder::OpenEncoder(const EncoderParam &encoderParam)
{
	CloseEncoder();

	if (m_pH264Buf)
	{
		delete [] m_pH264Buf;
		m_pH264Buf = NULL;
	}

	m_nSrcWidth = encoderParam.nWidth;
	m_nSrcHeight = encoderParam.nHeight;	
	
	m_hEncodeHandle = H264enc_open(encoderParam.nWidth, encoderParam.nHeight, encoderParam.nFrame, encoderParam.nQuant, encoderParam.nBitRate, encoderParam.nQMin, encoderParam.nQMax, encoderParam.nThreadNum);
	if (m_hEncodeHandle != NULL)
	{
		m_param = encoderParam;

		m_pFrameRGB = avcodec_alloc_frame();
		int nRGBNum = avpicture_get_size(PIX_FMT_RGB32, m_nSrcWidth, m_nSrcHeight);
		m_pRGBBuff = (uint8_t *)av_malloc(nRGBNum*sizeof(uint8_t));
		avpicture_fill((AVPicture *)m_pFrameRGB, m_pRGBBuff, PIX_FMT_RGB32, m_nSrcWidth, m_nSrcHeight);

		m_pFrameYUV = avcodec_alloc_frame();
		int nYUVNum = avpicture_get_size(PIX_FMT_YUV420P, encoderParam.nWidth, encoderParam.nHeight);
		m_pYUVBuff = (uint8_t *)av_malloc(nYUVNum*sizeof(uint8_t));
		avpicture_fill((AVPicture *)m_pFrameYUV, m_pYUVBuff, PIX_FMT_YUV420P, encoderParam.nWidth, encoderParam.nHeight);

		m_pRGB_YUV_CTX = sws_getContext(m_nSrcWidth, m_nSrcHeight, PIX_FMT_RGB32,
			encoderParam.nWidth, encoderParam.nHeight, PIX_FMT_YUV420P,
			SWS_BICUBIC, NULL,NULL,NULL);

		m_nBufLen = encoderParam.nWidth * encoderParam.nHeight;
		m_pH264Buf = new char[m_nBufLen];
	}	

	return m_hEncodeHandle != NULL;
}

void CVideoEncoder::CloseEncoder(void)
{
	if (m_hEncodeHandle)
	{
		H264enc_close(m_hEncodeHandle);
		m_hEncodeHandle = NULL;
	}

	if (m_pRGB_YUV_CTX)
	{
		sws_freeContext(m_pRGB_YUV_CTX);
		m_pRGB_YUV_CTX = NULL;
	}
	
	if (m_pRGBBuff)
	{
		av_free(m_pRGBBuff);
		m_pRGBBuff = NULL;
	}
	
	if (m_pFrameRGB)
	{
		av_free(m_pFrameRGB);
		m_pFrameRGB = NULL;
	}

	if (m_pFrameYUV)
	{
		av_free(m_pFrameYUV);
		m_pFrameYUV = NULL;
	}

	if (m_pYUVBuff)
	{
		av_free(m_pYUVBuff);
		m_pYUVBuff = NULL;
	}
}

int CVideoEncoder::Encode(RGBDataBuff *pRGB)
{
	if (pRGB->nWidth != m_nSrcWidth || pRGB->nHeigth != m_nSrcHeight)
	{
		CloseEncoder();

		m_nSrcWidth = pRGB->nWidth;
		m_nSrcHeight = pRGB->nHeigth;

		if (!OpenEncoder(m_param))
			return -1;
	}

	if (m_hEncodeHandle == NULL)
	{
		OpenEncoder(m_param);
		if (m_hEncodeHandle == NULL)
			return -1;
	}

	int nKeyFrame = 0;

	m_pFrameRGB->data[0] = (uint8_t *)pRGB->pData;

	sws_scale(m_pRGB_YUV_CTX, m_pFrameRGB->data, m_pFrameRGB->linesize, 0,
		pRGB->nHeigth, m_pFrameYUV->data, m_pFrameYUV->linesize);
	

	int nCurrentLen = 0;
	unsigned char *pYData = (unsigned char *)m_pFrameYUV->data[0];

	nCurrentLen += m_nSrcWidth * m_nSrcHeight;
	unsigned char *pUData = pYData + nCurrentLen;

	nCurrentLen += m_nSrcWidth * m_nSrcHeight / 4;
	unsigned char *pVData = pYData + nCurrentLen;

	unsigned char *pYV12[3];
	pYV12[0] = (unsigned char*)pYData;
	pYV12[1] = (unsigned char*)pUData;
	pYV12[2] = (unsigned char*)pVData;

	int iRet = H264enc_encode(m_hEncodeHandle, pYV12, (unsigned char*)m_pH264Buf, m_nBufLen, &nKeyFrame);
	if (iRet > 0 && m_pNotify)
	{
		H264VideoParam h264VideoFrame;
		h264VideoFrame.pBuf = m_pH264Buf;
		h264VideoFrame.nDataLen = iRet;
		h264VideoFrame.nWidth = m_param.nWidth;
		h264VideoFrame.nHeight = m_param.nHeight;
		h264VideoFrame.nFrameRate = m_param.nFrame;
		h264VideoFrame.nIsKeyFrame = nKeyFrame;
		h264VideoFrame.ulTimestamp = pRGB->ulTimestamp;

		m_pNotify->OnEncodeData(&h264VideoFrame);

		memset(m_pH264Buf, 0x00, m_nBufLen);
	}

	return iRet;
}

void CVideoEncoder::ThreadProcMain(void)
{
	while (!m_bExit)
	{
		while (m_lstRGBData.size() && m_bEncode)
		{
			if (m_bExit)
				return;

			RGBDataBuff *pRGB = NULL;

			{
				CAutoLock autoLock(&m_lockRGBData);
				pRGB = (RGBDataBuff *)m_lstRGBData.front();
				m_lstRGBData.pop_front();
			}

			if (pRGB)
				Encode(pRGB);

			ClearRGBBuf(pRGB);

			if (m_bExit)
				return;
		}

		my_sleep(5);
	}
}
